#include "llvm/ADT/Statistic.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/Function.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Demangle/Demangle.h"

using namespace llvm;

#define DEBUG_TYPE "test"

STATISTIC(test_counter, "Counts number of functions greeted");

int basic_block_handle(Function &F) {

	for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
		BasicBlock &Block = *BB;
		errs() << "BasicBlock: " << Block.getName() << "\n";
		for (BasicBlock::iterator I = Block.begin(), E = Block.end(); I != E; ++I) {
			Instruction &Inst = *I;

			// print instructions
			Inst.print(errs());
			errs() << "\n";
		}
	}
	return 0;
}

// Test Pass
namespace {
	struct Test: public FunctionPass {
		static char ID;
		Test() : FunctionPass(ID) {}

		bool runOnFunction(Function &F) override {
			++test_counter;
			errs() << "test: ";
			errs().write_escaped(F.getName()) << '\n';

			basic_block_handle(F);

			return false;
		}

		// We don't modify the program, so we preserve all analyses.
		void getAnalysisUsage(AnalysisUsage &AU) const override {
			AU.setPreservesAll();
		}
	};
}

char Test::ID = 0;
static RegisterPass<Test> A("test", "test pass");

// Module Pass
namespace {
	struct MyModulePass: public ModulePass {
		static char ID; // Pass identification, replacement for typeid
		MyModulePass() : ModulePass(ID) {}

		bool runOnModule(Module &M) override {
			for (Module::iterator F = M.begin(), E = M.end(); F != E; ++F) {
				Function &Func = *F;
				auto name = Func.getName();
				if(name.startswith("__Z") || name.startswith("_Z")) {
					std::string str = name.str();
					std::string demangled = demangle(str);
					errs() << "Function: " << demangled << "\n";
				} else {
					errs() << "Function: " << name.str() << "\n";
				}
			}
			return false;
		}
	};
}

char MyModulePass::ID = 0;
static RegisterPass<MyModulePass> B("test_module", "test pass for module", false, false);

// count the number of instructions in a function
void fun_count(Function &F) {
	unsigned InstructionCount = 0;

	for (Function::iterator BB = F.begin(), E = F.end(); BB != E; ++BB) {
		for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I) {
			InstructionCount++;
		}
	}

	errs() << "Function '" << F.getName() << "' has " << InstructionCount << " instructions.\n";
}

void fun_info(Function &F) {
	// get funtion's return type
	Type *retType = F.getReturnType();
	errs() << "Return Type: ";
	retType->print(errs());
	errs() << "\n";

	// dump fuction arguments
	if(F.arg_size() == 0) {
		errs() << "Have no arguments";
	} else {
		errs() << "Have " << F.arg_size() << " arguments, it's:\n";
	}

	for (Function::arg_iterator AI = F.arg_begin(), AE = F.arg_end(); AI != AE; ++AI) {
		Argument &arg = *AI;
		// errs() << "  Argument Type: " << arg.getType()->getName() << " ";
		errs() << "Number: " << arg.getArgNo() << ", Type: ";
		arg.print(errs());
		errs() << "\n";
	}

	errs() << "\n\n";
}

// My Function Pass
namespace {
	struct MyFunctionPass: public FunctionPass {
		static char ID;
		MyFunctionPass() : FunctionPass(ID) {}

		bool runOnFunction(Function &F) override {
			++test_counter;

			fun_count(F);
			fun_info(F);

			return false;
		}

		// We don't modify the program, so we preserve all analyses.
		void getAnalysisUsage(AnalysisUsage &AU) const override {
			AU.setPreservesAll();
		}
	};
}

char MyFunctionPass::ID = 0;
static RegisterPass<MyFunctionPass> C("test_function", "test pass for function");
